import torch
from torch import nn
from torch import optim
from tqdm import tqdm
import io


class Trainer:
    def __init__(self, criterion:nn.Module, optimizer:optim.Optimizer, model:nn.Module):
        self.criterion = criterion
        self.optimizer = optimizer
        self.model = model


    def train(self, trainloader, valloader, epoch, valepoch=3):
        best_val = float('inf')
        patience = 3
        patience_last = patience
        for i in range(epoch):
            total_loss = self.train_epoch(trainloader)
            print(f"training {i}th epoch, train loss is {total_loss / len(trainloader)}")
            if i % valepoch == 0:
                val_loss = self.val_epoch(valloader)
                print(f"val loss is {val_loss / len(valloader)}")
                if val_loss < best_val:
                    best_val = val_loss
                    patience_last = patience
                    buffer = io.BytesIO()
                    torch.save(self.model.state_dict(), buffer)
                    # 将指针重置到开头，准备加载
                    buffer.seek(0)
                else:
                    patience_last -= 1
                    print(f'early stop {patience - patience_last} / {patience}')
                    if patience_last == 0:
                        self.model.load_state_dict(torch.load(buffer))
                        break

    def train_epoch(self, dataloader):
        self.model.train()
        total_loss = 0
        for x, y in tqdm(dataloader):
            total_loss += self.train_step(x, y)
        return total_loss
    

    def train_step(self, x, y):
        self.optimizer.zero_grad()
        # with codebook
        # y_hat, z, hx = self.model(x)
        # loss = self.criterion(y, y_hat)
        # loss_codebook = self.criterion(z, hx)
        # loss_total = 0.7 * loss + 0.3 * loss_codebook
        # loss_total.backward()
        # without codebook
        y_hat = self.model(x)
        loss = self.criterion(y, y_hat)
        loss.backward()
        self.optimizer.step()
        return loss.item()
    

    def val_epoch(self, dataloader):
        self.model.eval()
        total_loss = 0
        for x, y in dataloader:
            total_loss += self.val_step(x, y)
        return total_loss
    

    def val_step(self, x, y):
        y_hat = self.model(x)
        loss = self.criterion(y, y_hat)
        return loss.item()
    

    def predict(self, dataloader, mean, std, target_col):
        self.model.eval()
        y_true = []
        y_pred = []
        
        for x, y in dataloader:
            y_t = y * std[target_col] + mean[target_col]
            y_h = self.model(x) * std[target_col] + mean[target_col]
            y_true += y_t.tolist()
            y_pred += y_h.tolist()

        return y_pred, y_true